import numpy as np
from os.path import join
from torch.utils.data import Dataset
from PIL import Image
from data.vfitransforms import rand_crop, rand_flip
import os


class Gazebo(Dataset):
    def __init__(
        self,
        db_dir,
        channels,
        augment_s=True,
        augment_t=True,
    ):
        self.channels = channels
        self.augment_s = augment_s
        self.augment_t = augment_t
        self.db_dir = db_dir
        self.data = self._prepare(db_dir)

    def __len__(self):
        return len(self.data)

    def _prepare(self, db_dir):
        raise NotImplementedError

    def __getitem__(self, index):
        gt_path, raw_path = self.data[index]
        gt = Image.open(join(self.db_dir, gt_path))
        raw = Image.open(join(self.db_dir, raw_path))
        if self.augment_s:
            gt, raw = rand_flip(gt, raw, p=0.5)
        if self.augment_t:
            gt, raw = rand_crop(gt, raw, sz=(256, 256))
        gt = np.array(gt, dtype=np.float32).squeeze()
        raw = np.array(raw, dtype=np.float32).squeeze()
        return {"raw": raw / 127.5 - 1.0, "gt": gt / 127.5 - 1.0}


class GazeboTrain(Gazebo):

    def _prepare(self, db_dir):
        data = []
        with open(join(db_dir, "train.txt")) as f:
            for gt, raw in map(lambda x: x.strip().split(","), f):
                data.append((gt, raw))
        return data


class GazeboValidate(Gazebo):

    def _prepare(self, db_dir):
        data = []
        with open(join(db_dir, "val.txt")) as f:
            for gt, raw in map(lambda x: x.strip().split(","), f):
                data.append((gt, raw))
        return data


class GazeboTest(Gazebo):

    def _prepare(self, db_dir):
        data = []
        with open(join(db_dir, "val.txt")) as f:
            for gt, raw in map(lambda x: x.strip().split(","), f):
                data.append((gt, raw))
        with open(join(db_dir, "train.txt")) as f:
            for gt, raw in map(lambda x: x.strip().split(","), f):
                data.append((gt, raw))
        return data


if __name__ == "__main__":
    dataset = GazeboTest(
        db_dir="/disk527/sdb1/a804_cbf/datasets/collect_data",
        channels=3,
        augment_s=False,
        augment_t=False,
    )
    for i, data in enumerate(dataset):
        print(data["raw"].shape)
        print(data["gt"].shape)
